import os
import sys
import time
import random
import math
import cv2
import numpy as np

# global variables for color
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (0, 0, 255)           # left
GREEN = (0, 255, 0)         # right
BLUE = (255, 0, 0)          # down
GRAY = (128, 128, 128)
DEEP_GRAY = (64, 64, 64)
BROWN = (0, 64, 128)        # up
LIGHT_GRAY = (192, 192, 192)
# create an numpy array of colors
colors = [BLACK, GREEN, RED, BLUE, GRAY, DEEP_GRAY, BROWN]

tile_size = 9

# class Map
class Map:
    """
    Map class
    """
    def __init__(self, width=tile_size, height=tile_size):
        """
        Constructor
        """
        self.width = width
        self.height = height
        self.map = np.zeros((height, width), np.uint8)
        self.pos = (0, 0)

    def set_tile(self, pos, tile):
        """
        Set tile at pos
        """
        # check if pos is in map
        if pos[0] < 0 or pos[0] >= self.width or pos[1] < 0 or pos[1] >= self.height:
            return 
        self.map[pos[1]][pos[0]] = tile

    def get_tile(self, pos):
        """
        Get tile at pos
        """
        return self.map[pos[1]][pos[0]]

    def reset(self, width=tile_size, height=tile_size):
        """
        Reset map
        """
        self.map = np.zeros((self.height, self.width), np.uint8)
    # save the map to a PNG image of size width x height
    def save(self, filename, width=tile_size, height=tile_size):
        """
        Save map to filename
        """
        img = np.zeros((height, width, 3), np.uint8)
        for y in range(height):
            for x in range(width):
                img[y][x] = colors[self.map[y][x]]
        cv2.imwrite(filename, img)

# mouse position
pos = (0, 0)
# create the map
map = Map()
# create the image
image = np.zeros((800, 800, 3), np.uint8)
# Get the image size
height, width, channels = image.shape
# tile color
tile_color = 0
# tiles
tile_x = 0
tile_y = 0
l_pressed = False

# mouse callback function
def mouse_cb(event, x, y, flags, param):
    """
    Mouse callback function
    """
    global pos
    global tile_color
    global map
    global height, width, channels
    global tile_x, tile_y
    global l_pressed
    # if left button is pressed
    if event == cv2.EVENT_LBUTTONUP:
        # set the global variable pos to the mouse position
        pos = (x, y)
        # clip the mouse position to the image size
        if pos[0] < 0:
            pos = (0, pos[1])
        if pos[1] < 0:
            pos = (pos[0], 0)
        if pos[0] > width-1:
            pos = (width-1, pos[1])
        if pos[1] > height-1:
            pos = (pos[0], height-1)
        # find out which tile the mouse is over
        tile_x = math.floor(pos[0] / (width/tile_size))
        tile_y = math.floor(pos[1] / (height/tile_size))
        # set the tile
        map.set_tile((tile_x, tile_y), tile_color)
        l_pressed = False
    # if left button is pressed
    if event == cv2.EVENT_LBUTTONDOWN:
        l_pressed = True
    # mouse move
    elif event == cv2.EVENT_MOUSEMOVE:
        pos = (x, y)
        # find out which tile the mouse is over
        tile_x = math.floor(pos[0] / (width/tile_size))
        tile_y = math.floor(pos[1] / (height/tile_size))
        if l_pressed:
            # find out which tile the mouse is over
            tile_x = math.floor(pos[0] / (width/tile_size))
            tile_y = math.floor(pos[1] / (height/tile_size))
            # set the tile
            map.set_tile((tile_x, tile_y), tile_color)

def main():
    """
    Main function
    """

    global pos
    global tile_color
    global map
    global height, width, channels
    global colors
    global image
    global tile_x, tile_y

    # create the window
    cv2.namedWindow("Original")
    # set the mouse callback function
    cv2.setMouseCallback("Original", mouse_cb)

    # main loop
    while True:

        image1 = image.copy()

        # draw the map on the image
        for y in range(map.height):
            for x in range(map.width):
                tile = map.get_tile((x, y))
                cv2.rectangle(image1, (int(x * (width/tile_size)), int(y * (height/tile_size))),
                              (int((x + 1) * (width/tile_size)), int((y + 1) * (height/tile_size))),
                              colors[tile], -1)
                # draw arrow on the tile
                if tile == 2:
                    # left arrow
                    cv2.arrowedLine(image1, (int((x + 0.75) * (width/tile_size)), int((y + 0.5) * (height/tile_size))),
                                    (int((x + 0.25) * (width/tile_size)), int((y + 0.5) * (height/tile_size))),
                                    WHITE, 3)
                elif tile == 1:
                    # right arrow
                    cv2.arrowedLine(image1, (int((x + 0.25) * (width/tile_size)), int((y + 0.5) * (height/tile_size))),
                                    (int((x + 0.75) * (width/tile_size)), int((y + 0.5) * (height/tile_size))),
                                    WHITE, 3)
                elif tile == 3:
                    # down arrow
                    cv2.arrowedLine(image1, (int((x + 0.5) * (width/tile_size)), int((y + 0.25) * (height/tile_size))),
                                    (int((x + 0.5) * (width/tile_size)), int((y + 0.75) * (height/tile_size))),
                                    WHITE, 3)
                elif tile == 6:
                    # up arrow
                    cv2.arrowedLine(image1, (int((x + 0.5) * (width/tile_size)), int((y + 0.75) * (height/tile_size))),
                                    (int((x + 0.5) * (width/tile_size)), int((y + 0.25) * (height/tile_size))),
                                    WHITE, 3)

        # draw a 8x8 grid on the image
        for i in range(0, width, int(width/tile_size)):
            cv2.line(image1, (i, 0), (i, height), (255, 255, 255))
        for i in range(0, height, int(height/tile_size)):
            cv2.line(image1, (0, i), (width, i), (255, 255, 255))

        # draw a circle of tile_color at left top corner of the image
        cv2.circle(image1, (50, 50), 11, WHITE, -1)
        cv2.circle(image1, (50, 50), 10, colors[tile_color], -1)
        # draw an arrow side by side with the tile_color circle
        if tile_color == 2:
            # left arrow
            cv2.arrowedLine(image1, (50, 50), (50 - 20, 50), WHITE, 3)
        elif tile_color == 1:
            # right arrow
            cv2.arrowedLine(image1, (50, 50), (50 + 20, 50), WHITE, 3)
        elif tile_color == 3:
            # down arrow
            cv2.arrowedLine(image1, (50, 50), (50, 50 + 20), WHITE, 3)
        elif tile_color == 6:
            # up arrow
            cv2.arrowedLine(image1, (50, 50), (50, 50 - 20), WHITE, 3)

        
        # highlight the tile the mouse is over
        cv2.rectangle(image1, (int(tile_x * (width/tile_size)), int(tile_y * (height/tile_size))),
                              (int((tile_x + 1) * (width/tile_size)), int((tile_y + 1) * (height/tile_size))),
                              RED, 1)

        cv2.imshow("Original", image1)
        # get key input
        key = cv2.waitKey(1)
        # press up key to change the tile color
        if key == ord('w'):
            tile_color = 6
        # press down key to change the tile color
        elif key == ord('s'):
            tile_color = 3
        # press left key to change the tile color
        elif key == ord('a'):
            tile_color = 2
        # press right key to change the tile color
        elif key == ord('d'):
            tile_color = 1
        # set tile_color to gray if key == ord('1')
        elif key == ord('1'):
            tile_color = 4
        # set tile_color to deep gray if key == ord('3')
        elif key == ord('2'):
            tile_color = 5
        # exit if esc key is pressed
        elif key == 27:
            break
        # press 0 to reset the map
        elif key == ord('0'):
            map.reset()
        # press s to save the map
        elif key == ord('o'):
            map.save("map.png")


if __name__ == '__main__':
    main()